Do you think we could align this to the API of the unified flash attention funciton for V1?

e.g. in v1/attention/backends/flash_attention

flash_attn_varlen_func(
                q=query[:num_actual_tokens],
                k=key_cache,
                v=value_cache,
                out=output[:num_actual_tokens],
                cu_seqlens_q=attn_metadata.query_start_loc,
                max_seqlen_q=attn_metadata.max_query_len,
                cu_seqlens_k=attn_metadata.seq_start_loc,
                max_seqlen_k=attn_metadata.max_seq_len,
                softmax_scale=self.scale,
                causal=True,
                alibi_slopes=self.alibi_slopes,
                window_size=self.sliding_window,
                block_table=attn_metadata.block_table,
                softcap=self.logits_soft_cap,
            )

Aligning to this interface will reduce special cases.

Great question!

There are three challenges that are blocking the alignment of the interface:

1. calling reshape_and_cache_flash before flash_attn_varlen_func can be inefficient read-after-write, due to synchronous process on a heterogeneous architecture. On neuron stack, it's better to merge cached tokens with new tokens in an asynchronous fashion. Therefore, we pass both cached KV and active KV to the flash-attention function call, so that writing to HBM won't block compuation.
2. Since slicing num_actual_tokens would not reduce computation or bandwidth utilization, I think it's better to slice the logits in the last attention layer.
3. We find it more efficient to compuate attention mask before each of the layer. Therefore, instead of passing sequence lengths and query lengths, we pass the pre-computed attention mask directly to the flash-attention kernel.

I'm actively looking for ideas that can help close the gap between the interfaces, without degrading performance of the kernel.

@robertgshaw2-redhat do you think the interface misalignment can be a blocking issue for merging?

Here is a typo arthimetic breaks the CI.